package com.demo.server.interceptor;

import java.lang.reflect.Method;

import org.apache.commons.lang3.StringUtils;
import org.apache.ibatis.cache.CacheKey;
import org.apache.ibatis.executor.Executor;
import org.apache.ibatis.mapping.BoundSql;
import org.apache.ibatis.mapping.MappedStatement;
import org.apache.ibatis.mapping.SqlCommandType;
import org.apache.ibatis.mapping.SqlSource;
import org.apache.ibatis.plugin.Interceptor;
import org.apache.ibatis.plugin.Intercepts;
import org.apache.ibatis.plugin.Invocation;
import org.apache.ibatis.plugin.Plugin;
import org.apache.ibatis.plugin.Signature;
import org.apache.ibatis.session.ResultHandler;
import org.apache.ibatis.session.RowBounds;

import com.demo.server.annotation.SqlLimit;
import com.demo.server.constant.CommonConstant;
import com.demo.server.entity.User;
import com.demo.server.util.AccessTokenUtil;

import lombok.extern.slf4j.Slf4j;

/**
 * @title 数据权限
 * @date 2024-02-22
 * @author wulianwei
 */
@Slf4j
@Intercepts({
    @Signature(
            type = Executor.class,
            method = "query",
            args = {MappedStatement.class, Object.class, RowBounds.class, ResultHandler.class}
    ),
    @Signature(
            method = "query",
            type = Executor.class,
            args = {MappedStatement.class, Object.class, RowBounds.class, ResultHandler.class, CacheKey.class, BoundSql.class}
    ),
    @Signature(
            type = Executor.class,
            method = "update",
            args = {MappedStatement.class, Object.class}
    )
})
public class MybatisInterceptor implements Interceptor {

    @Override
    public Object intercept(Invocation invocation) throws Throwable {    	
        MappedStatement statement = (MappedStatement) invocation.getArgs()[0];
        Object parameter = invocation.getArgs()[1];
        BoundSql boundSql = statement.getBoundSql(parameter);
        String originalSql = boundSql.getSql();
        Object parameterObject = boundSql.getParameterObject();
        
        SqlCommandType sqlCommandType = statement.getSqlCommandType();

        if (SqlCommandType.INSERT.equals(sqlCommandType) || SqlCommandType.FLUSH.equals(sqlCommandType)) {
            return invocation.proceed();
        }
        
        SqlLimit sqlLimit = isLimit(statement);
        if (sqlLimit == null || sqlLimit.isWork() == false) {
            return invocation.proceed();
        }
        User user = AccessTokenUtil.getUser();
 	    if(user == null) {
 	    	return invocation.proceed();
 	    }
 	    if(user.getRoleIds() != null && user.getRoleIds().contains(CommonConstant.ROLE_MANAGER)) {
 	    	return invocation.proceed();
 	    }
        String limitValue = user.getId();
        String sql = addTenantCondition(originalSql, limitValue, sqlLimit.alis(),sqlLimit.columnName());
        log.info("原SQL：{}， 数据权限替换后的SQL：{}", originalSql, sql);
        BoundSql newBoundSql = new BoundSql(statement.getConfiguration(), sql, boundSql.getParameterMappings(), parameterObject);
        MappedStatement newStatement = copyFromMappedStatement(statement, new BoundSqlSqlSource(newBoundSql));
        invocation.getArgs()[0] = newStatement;
        return invocation.proceed();
    }

    @Override
    public Object plugin(Object target) {
        if (target instanceof Executor) {
            return Plugin.wrap(target, this);
        } else {
            return target;
        }
    }

	
    /**
     * 重新拼接SQL
     */
    private String addTenantCondition(String originalSql, String limitValue, String alias,String column) {
        String field = "tenant_id";
        if(StringUtils.isNoneBlank(column)){
            field = column;
        }
        if(StringUtils.isNoneBlank(alias)){
            field = alias + "." + field;
        }

        StringBuilder sb = new StringBuilder(originalSql);
        limitValue = "'"+limitValue+"'";
        int index = StringUtils.lastIndexOfAny(sb, "where","WHERE");
        if (index < 0) {
        	int orderIndex = sb.indexOf("order by");
        	if(orderIndex < 0) {
        		sb.append(" where ") .append(field).append(" = ").append(limitValue);
        	}else {
        		 sb.insert(orderIndex - 1, " where " + field +" = " + limitValue);
        	}
        } else {
            sb.insert(index + 5, " " + field +" = " + limitValue + " and ");
        }
        return sb.toString();
    }

    private MappedStatement copyFromMappedStatement(MappedStatement ms, SqlSource newSqlSource) {
        MappedStatement.Builder builder = new MappedStatement.Builder(ms.getConfiguration(), ms.getId(), newSqlSource, ms.getSqlCommandType());
        builder.resource(ms.getResource());
        builder.fetchSize(ms.getFetchSize());
        builder.statementType(ms.getStatementType());
        builder.keyGenerator(ms.getKeyGenerator());
        builder.timeout(ms.getTimeout());
        builder.parameterMap(ms.getParameterMap());
        builder.resultMaps(ms.getResultMaps());
        builder.cache(ms.getCache());
        builder.useCache(ms.isUseCache());
        return builder.build();
    }


    /**
     * 通过注解判断是否需要限制数据
     * @return
     */
    private SqlLimit isLimit(MappedStatement mappedStatement) {
        SqlLimit sqlLimit = null;
        try {
            String id = mappedStatement.getId();
            String className = id.substring(0, id.lastIndexOf("."));
            String methodName = id.substring(id.lastIndexOf(".") + 1, id.length());
            final Class<?> cls = Class.forName(className);
            if(cls.isAnnotationPresent(SqlLimit.class)) {
            	sqlLimit = cls.getAnnotation(SqlLimit.class);
            }
            final Method[] method = cls.getMethods();
            for (Method me : method) {
                if (me.getName().equals(methodName) && me.isAnnotationPresent(SqlLimit.class)) {
                    sqlLimit = me.getAnnotation(SqlLimit.class);
                }
            }
        } catch (Exception e) {
            e.printStackTrace();
        }
        return sqlLimit;
    }


    public static class BoundSqlSqlSource implements SqlSource {

        private final BoundSql boundSql;

        public BoundSqlSqlSource(BoundSql boundSql) {
            this.boundSql = boundSql;
        }

        @Override
        public BoundSql getBoundSql(Object parameterObject) {
            return boundSql;
        }
    }
}
